import numpy as np
from sklearn.linear_model import LinearRegression
# 升维模块
from sklearn.preprocessing import PolynomialFeatures
import matplotlib.pyplot as plt
# 计算mse损失函数，评估效果
from sklearn.metrics import mean_squared_error

np.random.seed(30)
m = 100
# np.random.rand(100,1)~[0,1],6*np.random.rand(100,1)~[0,6],6*np.random.rand(100,1)-3~[-3,3]
X = 6*np.random.rand(100,1)-3
Y = 0.5*X**2 + X + 2 + np.random.randn(100,1)
plt.plot(X,Y,'b.')

X_train = X[:80]
Y_train = Y[:80]
X_test = X[80:]
Y_test = Y[80:]

# 1维绿色直线，二维红色加号，10维黄色星号
d = {1: 'g-', 2: 'r+', 10: "y*"}
# for循环遍历字典时，i是键值对的键
for i in d:
    # degree 升到几维 include_bias：创建截距项
    poly_features = PolynomialFeatures(degree=i,include_bias=True)
    X_poly_train = poly_features.fit_transform(X_train)
    X_poly_test = poly_features.fit_transform(X_test)
    print(X_train[0])
    print(X_poly_train[0])
    print(X_train.shape)
    print(X_poly_train.shape)

    # 对训练集训练，拿测试集测试并计算MSE
    linear_reg = LinearRegression(fit_intercept=False)
    linear_reg.fit(X_poly_train,Y_train)
    print(f"linear_reg{i}_intercept:",linear_reg.intercept_)
    print(f"linear_reg{i}_coef:", linear_reg.coef_)
    print(f"linear_reg{i}_train_predict:", linear_reg.predict(X_poly_train))
    print(f"linear_reg{i}_test_predict:", linear_reg.predict(X_poly_test))
    # 看看是否随着degree的升维，是否过拟合了
    y_train_predict = linear_reg.predict(X_poly_train)
    y_test_predict = linear_reg.predict(X_poly_test)

    plt.plot(X_poly_train[:,1],y_train_predict,d[i])
    # 求解MSE损失函数
    print(f"degree{i}_mse_train:",mean_squared_error(Y_train,y_train_predict))
    print(f"degree{i}_mse_test:", mean_squared_error(Y_test, y_test_predict))

plt.show()